package com.zhuyuan.database.interceptor;

import com.baomidou.dynamic.datasource.toolkit.DynamicDataSourceContextHolder;
import org.apache.ibatis.executor.Executor;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.mapping.SqlCommandType;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.session.ResultHandler;
import org.apache.ibatis.session.RowBounds;
import org.springframework.stereotype.Component;

import java.util.Properties;

/**
 * @Author: 张琳凯
 * @Description: 通过 MyBatis 拦截器解析 SQL 类型，自动切换数据源,写读分离
 * @DateTime: 2025/2/27 19:16
 **/
@Intercepts({
        @Signature(type = Executor.class, method = "update", args = {MappedStatement.class, Object.class}),
        @Signature(type = Executor.class, method = "query", args = {MappedStatement.class, Object.class, RowBounds.class, ResultHandler.class})
})
@Component
public class ReadWriteInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        MappedStatement ms = (MappedStatement) invocation.getArgs()[0];
        SqlCommandType sqlCommandType = ms.getSqlCommandType();

        // 写操作（INSERT/UPDATE/DELETE）走主库
        if (SqlCommandType.INSERT == sqlCommandType ||
                SqlCommandType.UPDATE == sqlCommandType ||
                SqlCommandType.DELETE == sqlCommandType) {
            //DynamicDataSourceContextHolder.push("master");
        }
        // 读操作（SELECT）走从库
        else if (SqlCommandType.SELECT == sqlCommandType) {
            //DynamicDataSourceContextHolder.push("slave");
        }

        try {
            return invocation.proceed();
        } finally {
            DynamicDataSourceContextHolder.clear();
        }
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {}
}
